#!/usr/bin/env python3
"""
DataPlexus 僵尸任务清理工具
清理长时间运行但实际已停止的同步任务
"""

import pymysql
import sys
from datetime import datetime, timedelta

# 数据库配置
DB_CONFIG = {
    'host': '192.168.37.11',
    'port': 3306,
    'user': 'root',
    'password': '123456',
    'database': 'data_plexus',
    'charset': 'utf8mb4'
}

def cleanup_zombie_tasks(max_runtime_hours=2):
    """
    清理僵尸任务
    
    Args:
        max_runtime_hours: 最大运行时间（小时），超过此时间的运行任务将被视为僵尸任务
    """
    connection = None
    try:
        print("DataPlexus 僵尸任务清理工具")
        print("=" * 50)
        
        # 连接数据库
        print("连接数据库...")
        connection = pymysql.connect(**DB_CONFIG)
        cursor = connection.cursor()
        
        print(f"成功连接到数据库: {DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}")
        
        # 计算截止时间
        cutoff_time = datetime.now() - timedelta(hours=max_runtime_hours)
        print(f"查找运行时间超过 {max_runtime_hours} 小时的任务（截止时间: {cutoff_time}）")
        
        # 查找僵尸任务
        cursor.execute("""
            SELECT 
                id,
                task_id,
                status,
                started_at,
                created_at,
                TIMESTAMPDIFF(MINUTE, started_at, NOW()) as runtime_minutes
            FROM dp_sync_executions 
            WHERE status IN ('running', 'paused') 
              AND started_at < %s
            ORDER BY started_at
        """, (cutoff_time,))
        
        zombie_tasks = cursor.fetchall()
        
        if not zombie_tasks:
            print("✅ 没有发现僵尸任务")
            return True
        
        print(f"🔍 发现 {len(zombie_tasks)} 个可能的僵尸任务:")
        print("执行ID\t任务ID\t状态\t\t开始时间\t\t\t运行时间(分钟)")
        print("-" * 80)
        
        for task in zombie_tasks:
            exec_id, task_id, status, started_at, created_at, runtime_minutes = task
            print(f"{exec_id}\t{task_id}\t{status}\t{started_at}\t{runtime_minutes}")
        
        # 询问是否清理
        print(f"\n是否要清理这 {len(zombie_tasks)} 个僵尸任务？(y/N): ", end="")
        response = input().strip().lower()
        
        if response not in ['y', 'yes']:
            print("取消清理操作")
            return True
        
        # 执行清理
        print("\n开始清理僵尸任务...")
        cleaned_count = 0
        
        for task in zombie_tasks:
            exec_id, task_id, status, started_at, created_at, runtime_minutes = task
            
            try:
                # 更新任务状态为失败
                cursor.execute("""
                    UPDATE dp_sync_executions 
                    SET 
                        status = 'failed',
                        completed_at = %s,
                        updated_at = %s,
                        error_message = %s
                    WHERE id = %s AND status IN ('running', 'paused')
                """, (
                    datetime.now(),
                    datetime.now(),
                    f'任务异常终止：运行时间超过{max_runtime_hours}小时，疑似僵尸进程',
                    exec_id
                ))
                
                if cursor.rowcount > 0:
                    print(f"✅ 清理执行ID {exec_id}（任务ID {task_id}）")
                    cleaned_count += 1
                else:
                    print(f"⚠️ 执行ID {exec_id} 可能已被其他进程处理")
                    
            except Exception as e:
                print(f"❌ 清理执行ID {exec_id} 失败: {e}")
        
        # 提交事务
        connection.commit()
        
        print(f"\n✅ 成功清理了 {cleaned_count} 个僵尸任务")
        
        # 显示清理后的状态
        if cleaned_count > 0:
            print("\n验证清理结果...")
            cursor.execute("""
                SELECT 
                    COUNT(*) as total_running,
                    COUNT(CASE WHEN started_at < %s THEN 1 END) as old_running
                FROM dp_sync_executions 
                WHERE status IN ('running', 'paused')
            """, (cutoff_time,))
            
            result = cursor.fetchone()
            total_running, old_running = result
            
            print(f"当前运行中的任务数: {total_running}")
            print(f"运行时间超过 {max_runtime_hours} 小时的任务数: {old_running}")
            
            if old_running == 0:
                print("✅ 所有僵尸任务已清理完成!")
            else:
                print(f"⚠️ 仍有 {old_running} 个长时间运行的任务")
        
        return True
        
    except Exception as e:
        print(f"清理失败: {e}")
        if connection:
            connection.rollback()
        return False
        
    finally:
        if connection:
            connection.close()
            print("\n数据库连接已关闭")

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='DataPlexus 僵尸任务清理工具')
    parser.add_argument(
        '--max-hours', 
        type=int, 
        default=2, 
        help='最大运行时间（小时），默认2小时'
    )
    parser.add_argument(
        '--auto', 
        action='store_true', 
        help='自动清理，不询问确认'
    )
    
    args = parser.parse_args()
    
    # 如果是自动模式，修改input函数
    if args.auto:
        global input
        input = lambda prompt: 'y'
        print("🤖 自动清理模式")
    
    success = cleanup_zombie_tasks(args.max_hours)
    
    if success:
        print("\n✓ 清理完成!")
        sys.exit(0)
    else:
        print("\n✗ 清理失败!")
        sys.exit(1)

if __name__ == "__main__":
    main()
